#include "list_sieve.h"

ZZ shortest_mag_sq;

int main(int argc, char* argv[])
{
  Sampler S;
  for(int d = atoi(argv[1]); d <= atoi(argv[2]); d++)
  {
    for(int i = 0; i < atoi(argv[3]); i++)
    {
      mat_ZZ B = S.randomBasis(d);
      ListSieve solver(B);
      vec svp = solver.g_solve();
      vec svp2 = solver.solve();

      mat_ZZ LLL_basis = B;
      ZZ det;
      long i = LLL(det, LLL_basis);
      shortest_mag_sq = mag_sq(LLL_basis[LLL_basis.NumRows() - i]);

      cout << (shortest_mag_sq >= svp.mag_sq) << endl;
      cout << (shortest_mag_sq >= svp2.mag_sq) << endl;
    }
  }

  return 0;
}

/* 
 * Constructor for ListSieve, initializes variables that we may need later
 */
ListSieve::ListSieve(mat_ZZ basis)
{
  B = basis;
  G_BKZ_FP(B, 0.99, 20);
  dimension = B.NumRows();
  srand(time(0));

  //this is used for sampleZ
  ListSieve::t_ = log(dimension);
  ListSieve::pi = atan(1)*4;
}

/*
 * Solves SVP by generating random lattice vectors(using 0 and 1 only) and
 * then reducing them using lattice vectors from a list until we can no
 * longer find a shorter vector
 */
vec ListSieve::solve()
{
  vector<vec> list;
  int collisions = 0;
  int max_collisions = dimension;
  int iterations = 0;
  vec t;
  //while we haven't hit threshold number of collisions
  while(collisions < max_collisions)
  {
    //get a random lattice vector
    t = randomVec(B);
    //reduce it using vectors from initially empty list
    list_reduce(&t, list);
    if(t.mag_sq == 0)
      //if we get a collision increment collisions
      collisions++;
    else
      //if not a collision we add it to the list, keeping the list sorted
      sorted_add(&t, list);
    iterations++;
  }
  cout << "iterations: " << iterations << endl;
  cout << "size: " << list.size() << endl;
  return list[0];
}


/* 
 * Returns a random lattice vector that can be obtained by taking the sum
 * of some subset of the basis vectors
 */
vec ListSieve::randomVec(mat_ZZ &basis)
{
  vector<long> v;//0 1 vector
  vec ret;
  ZZ zero;
  zero = 0;

  v.resize(basis.NumCols(), 0);
  //set v to all 0 or 1s randomly
  for(int i = 0; i < basis.NumCols(); i++)
  {
    v[i] = rand()%2;
  }
  //initialize ret to all 0s
  for(int i = 0; i < basis.NumCols(); i++)
    append(ret.v, zero);
  for(int i = 0; i < ret.v.length(); i++)
  {
    if(v.at(i)) //if it's 1 we have to add it to ret
      ret.v = ret.v + basis[i];
  }
  ret.mag_sq = mag_sq(ret.v);
  if(ret.mag_sq == 0)
    return randomVec(basis);

  return ret;
}

/*
 * Returns a random integer with probability assigned by a discrete gaussian
 * around c. 
 * Copied from GSieve and the paper
 */
long ListSieve::sampleZ(double c, double s_squared)
{
  double s = sqrt(s_squared);
  long minimum_c = floor(c - s * t_);
  long maximum_c = ceil(c + s * t_);
  long x;
  double rho;
  while(true) {
    x = minimum_c +
      round((maximum_c - minimum_c) * double(rand()) / RAND_MAX);
    rho = exp( - pi * (x - c) * (x - c) / s_squared);
    if ((double(rand()) / RAND_MAX) <= rho) {
        return x;
    }
  }
}

/*
 * reduces t using vectors in the list, which we assume to be sorted
 * We stop when we reach the end of the list and only use vectors that are
 * longer than t(as an optimization).
 * Note that if t is reduced we start from the beginning again
 */
void ListSieve::list_reduce(vec *t, vector<vec> &list)
{
  for(int i = 0; i < list.size(); i++)
  {
    //if the list vector is longer we are done!
    if(list[i].mag_sq > t->mag_sq)
      break;
    vec red = reduce(t, &list[i]);
    if(red.mag_sq < t->mag_sq)
    {
      *t = red;
      i = 0;
    }
  }
}

/* 
 * This function adds t to list while keeping it sorted, assuming it is
 * already sorted
 */
void ListSieve::sorted_add(vec *t, vector<vec> &list)
{
  //we start from the top since it's more likely that we will add a short one
  //just increments i till we are at a point that t is shorter than list[i]
  int i = 0;
  while(i < list.size() && list[i].mag_sq < t->mag_sq)
    i++;
  list.insert(list.begin() + i, *t);;
}

/*
 * Function adds t to list while keeping it sorted, unless t is already in
 * the list!
 */
void ListSieve::g_sorted_add(vec *t, vector<vec> &list)
{
  //we start from the top since it's more likely that we will add a short one
  //just increments i till we are at a point that t is shorter than list[i]
  int i = 0;
  while(i < list.size() && list[i].mag_sq < t->mag_sq)
    i++;
  //stop duplicates
  if(i == list.size() || list.size() == 0 || list[i].v != t->v)
    list.insert(list.begin() + i, *t);;
}


/*
 * Returns v - c * a where c is a random integer distributed according to a 
 * discrete gaussian around t | t is the int that minimizes v - c * a
 */
vec ListSieve::g_reduce(vec *v, vec *a)
{
  //first find c
  ZZ min; //integer that minimizes v - c * a
  ZZ dot; //dot product of v and a
  vec ret;

  //find c such that v - c* a is minimized
  InnerProduct(dot, v->v, a->v); //dot = inner product
  //c = round((t * x)/(x * x) ) where * is dot product
  RoundToZZ(min, MakeRR(dot, 0)/MakeRR(a->mag_sq,0));


  long center = to_long(min);
  long c = sampleZ(center, t_);

  ret.v = v->v - c * a->v;
  ret.mag_sq = v->mag_sq + c * c * a ->mag_sq - 2 * c * dot;
  return ret;
}

/* 
 * Does list sieve except using g_reduce instead of reduce. Note though that 
 * we might make the vector longer.
 * To account for this, we continue through the list if it doesn't get shorter,
 * (including if it gets longer), but we only start from the beginning again
 * if it gets shorter
 */
void ListSieve::g_list_reduce(vec *t, vector<vec> &list)
{
  for(int i = 0; i < list.size(); i++)
  {
    //if the list vector is longer we are done!
    if(list[i].mag_sq > t->mag_sq)
      break;
    vec red = g_reduce(t, &list[i]);
    if(red.mag_sq < t->mag_sq)
    {
      i = 0;
    }
    *t = red;
  }
}


/*
 * Same as solve except that it uses g_reduce instead of reduce.
 * However since collisions are harder to find, we continue until we hit some
 * max of times that we don't find a shorter vector
 */
vec ListSieve::g_solve()
{
  vector<vec> list;
  int collisions = 0;
  int max_collisions = dimension/2;
  vec t;
  int iterations = 0;

  //while we haven't hit threshold number of collisions
  while(collisions < max_collisions)
  {
    iterations++;
    //get a random lattice vector
    t = randomVec(B);
    //reduce it using vectors from initially empty list
    g_list_reduce(&t, list);
    //first make sure it's not 0 before adding it to the list
    if(t.mag_sq != 0)
      sorted_add(&t, list);
    else
      collisions++;
      /*
    //remove long vectors
    for(int i = list.size()-1; list[i].mag_sq > list[0].mag_sq*sqrt(dimension);
        i--)
    {
      list.erase(list.end());
    }
    */
  }
  cout << "iterations: " << iterations << endl;
  cout << "size: " << list.size() << endl;

  return list[0];
}

